import os
import json
import sqlite3
from dotenv import load_dotenv
from mcp.server.fastmcp import FastMCP
from pathlib import Path

# 加载环境变量
load_dotenv(override=True)

# 创建MCP
mcp = FastMCP("AgentSQLiteServer")

# 构建文件读写路径
file_folder_path = Path(__file__).resolve().parent.parent.parent / "files"

# Tools
@mcp.tool()
async def sql_inquire(sql_query: str) -> str:
    """
    当用户需要进行数据库查询工作时，请调用该函数。
    该函数用于在指定SQLite数据库上运行一段SQL代码，完成数据查询相关工作，
    并且当前函数是使用Python标准库sqlite3连接SQLite数据库。
    本函数只负责运行SQL代码并进行数据查询。
    :param sql_query: 字符串形式的SQL查询语句，用于执行对SQLite数据库进行查询，并获得相关信息
    :return：sql_query在SQLite中的运行结果。
    """

    # 创建连接（SQLite使用文件路径，若文件不存在会自动创建）
    connection = sqlite3.connect(file_folder_path / 'dacp.db')
    # 设置行工厂，让结果包含字段名
    connection.row_factory = sqlite3.Row

    try:
        cursor = connection.cursor()
        cursor.execute(sql_query)
        # 获取所有查询结果
        results = cursor.fetchall()
        # 转换为包含字段名的字典列表
        formatted_results = [dict(row) for row in results]
    except Exception as e:
        # 捕获并返回错误信息
        return json.dumps({"error": str(e)}, ensure_ascii=False)
    finally:
        connection.close()

    # 将结果以JSON字符串形式返回
    return json.dumps(formatted_results, ensure_ascii=False)


if __name__ == "__main__":
    mcp.run(transport="stdio")